package com.xzx.masterslave.masterslave.config;

import com.xzx.masterslave.masterslave.contast.DataSouceType;
import com.xzx.masterslave.masterslave.context.DataSouceContextHolder;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import java.util.concurrent.atomic.AtomicInteger;

/**
 * 数据源切换规则
 * @author admin
 */
public class MyAbstrctRoutingDataSource extends AbstractRoutingDataSource {
    private final int dataSourceNumber;

    private AtomicInteger atomicInteger = new AtomicInteger(0);

    public MyAbstrctRoutingDataSource(int dataSourceNumber){
        this.dataSourceNumber = dataSourceNumber;
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String typeKey = DataSouceContextHolder.getJdbcType();
        if (DataSouceType.write.getType().equals(typeKey)){
            //如果当前ThreadLocal 保存类型 write, 返回 write
            return DataSouceType.write.getType();
        }

        //简单 读负载均衡
        int number = atomicInteger.getAndAdd(1);
        int lookUpKey = number % dataSourceNumber;
        return new Integer(lookUpKey);
    }
}
